"""
词频统计
@author:yczha
"""

from ast import keyword
from email import message
from errno import ENOEXEC
import os
import wx
from threading import Thread
from concurrent.futures import ThreadPoolExecutor, as_completed

from openpyxl import Workbook

from src.config import settings
from src.utils import get_allowed_files
from src.ui.transparent_stastictext import TransparentStaticText
from src.core.words_freq.words_freq_count import WordFreqCount


class WordsFreqCountUI(wx.Dialog):
    def __init__(self, parent):
        super(WordsFreqCountUI, self).__init__(parent, title="词频统计")
        self.keywords = []
        panel = wx.Panel(self)
        back = wx.StaticBitmap(
            panel, bitmap=wx.Bitmap(os.path.join(settings.assets_dir, "bg.png"))
        )
        self.result = None  # 存储结果

        tip = f"""*功能：PDF文件或txt文本文档词频统计，\n  原理：查找出现的关键词及其次数\n  *使用：1.选择关键词文件，一行一个关键词,\n                参考本程序附带的keywords.txt\n             2.选择要解析的文件夹"""
        self.tip = TransparentStaticText(
            self, label=tip, style=wx.TE_MULTILINE | wx.TE_RICH2
        )
        self.tip.SetFont(
            wx.Font(
                11,
                wx.FONTFAMILY_TELETYPE,
                wx.FONTSTYLE_ITALIC,
                wx.FONTWEIGHT_LIGHT,
                False,
                "Microsoft YaHei",
            )
        )
        self.tip.SetForegroundColour((255, 255, 255))
        self.tip.SetPosition((50, 180))

        select_btn = wx.BitmapButton(
            back,
            bitmap=wx.Bitmap(os.path.join(settings.assets_dir, "select-keywords.png")),
        )
        select_btn.SetPosition((100, 80))
        select_btn.Bind(wx.EVT_BUTTON, self.select_keywords)

        open_btn = wx.BitmapButton(
            back, bitmap=wx.Bitmap(os.path.join(settings.assets_dir, "open-file.png"))
        )
        open_btn.SetPosition((250, 80))
        open_btn.Bind(wx.EVT_BUTTON, self.open_file)

        panel.Fit()
        self.Fit()
        self.Centre()

    def open_file(self, e):
        dlg = wx.DirDialog(self, "选择文件夹", style=wx.DD_DEFAULT_STYLE)
        if dlg.ShowModal() == wx.ID_OK:
            allowed_files = get_allowed_files(
                dlg.GetPath(), suffix=[".pdf", ".PDF", ".txt"]
            )
            if len(allowed_files) < 1:
                wx.MessageBox("所选目录没有有效文件", "提示", wx.OK | wx.ICON_WARNING)
                return
            if not self.keywords:
                wx.MessageBox("请先选择关键词文件！", "警告", wx.OK | wx.ICON_WARNING)
                return
            task_queue = []
            with TaskProgressDialog(
                None, files=allowed_files, task_queue=task_queue, keywords=self.keywords
            ) as dlg:
                dlg.ShowModal()
            if len(task_queue) < 1:
                wx.MessageBox("结果为空", "提示", wx.OK | wx.ICON_WARNING)
            with wx.MessageDialog(
                None,
                "解析完成，保存结果?",
                "保存文件",
                wx.YES_NO | wx.YES_DEFAULT | wx.ICON_QUESTION,
            ) as dlg:
                if dlg.ShowModal() != wx.ID_YES:
                    return
            fdlg = wx.FileDialog(
                self,
                "保存结果",
                os.getcwd(),
                "词频统计.xlsx",
                "Excel file(*.xlsx)|*.xls",
                wx.FD_SAVE,
            )
            if fdlg.ShowModal() != wx.ID_OK:
                return
            state, msg = self.write_to_excel(task_queue, fdlg.GetPath())
            if state:
                wx.MessageBox(msg, "提示", wx.OK | wx.ICON_INFORMATION)
            else:
                wx.MessageBox(msg, "错误", wx.OK | wx.ICON_ERROR)

    def select_keywords(self, e):
        """选择关键词"""
        dlg = wx.FileDialog(
            self, "选择文件", defaultDir="./", defaultFile="keywords.txt", wildcard="*.txt"
        )
        keyword_path = ""
        if dlg.ShowModal() == wx.ID_OK:
            keyword_path = dlg.GetPath()
            if not os.path.isfile(keyword_path):
                wx.MessageBox("无效文件", "错误", wx.OK | wx.ICON_ERROR)
                return
        else:
            return
        try:
            self.keywords = WordFreqCount().get_keywords(keyword_path)
            if not self.keywords:
                wx.MessageBox("没有关键词！", "警告", wx.OK | wx.ICON_WARNING)
                return
            wx.MessageBox(
                f"已选择{len(self.keywords)}个关键词", "提示", wx.OK | wx.ICON_INFORMATION
            )
        except Exception as e:
            wx.MessageBox(f"读取关键词文件出错({e})", "错误提示", wx.OK | wx.ICON_ERROR)
            return

    def write_to_excel(self, task_data, path):
        wb = Workbook()
        ws = wb.active
        head_row = ["文件名称", "文件路径", "备注"] + self.keywords
        ws.append(head_row)
        for status, msg,file_path, words_cnt in task_data:
            filename, filepath = os.path.basename(file_path), os.path.dirname(file_path)
            if status == 1000:
                row_data = [filename, filepath, ""] + [
                    words_cnt[word] for word in self.keywords
                ]
            else:
                row_data = [filename, filepath, str(msg)] + [
                    "" for _ in self.keywords
                ]
            ws.append(row_data)
        try:
            wb.save(path)
            return True, "文件保存成功！"
        except Exception as e:
            return False, str(e)


class TaskProgressDialog(wx.Dialog):
    """进度条控制"""

    def __init__(self, parent, files, task_queue, keywords):
        super(TaskProgressDialog, self).__init__(parent, title="完成进度")

        self.task_queue = task_queue
        self.files = files
        self.keywords = keywords

        panel = wx.Panel(self)
        self.back = wx.StaticBitmap(
            panel, bitmap=wx.Bitmap(os.path.join(settings.assets_dir, "progress.png"))
        )
        self.STOP_TASK = False
        self.gauge = wx.Gauge(self.back, range=len(files), size=(320, 20))
        self.gauge.SetPosition((20, 50))
        self.file_lab = TransparentStaticText(self.back, label="0.0%")
        self.file_lab.SetFont(
            wx.Font(
                12,
                wx.FONTFAMILY_MODERN,
                wx.FONTSTYLE_NORMAL,
                wx.FONTWEIGHT_LIGHT,
                False,
                "Microsoft YaHei",
            )
        )
        self.file_lab.SetForegroundColour((0, 0, 255))
        self.file_lab.SetPosition((350, 50))
        self.file_msg = TransparentStaticText(self.back, label="")
        self.file_msg.SetFont(
            wx.Font(
                12,
                wx.FONTFAMILY_MODERN,
                wx.FONTSTYLE_NORMAL,
                wx.FONTWEIGHT_LIGHT,
                False,
                "Microsoft YaHei",
            )
        )
        self.file_msg.SetForegroundColour((255, 255, 255))
        self.file_msg.SetPosition((100, 100))

        panel.Fit()
        self.Fit()
        self.Centre()
        self.SetMaxSize(self.GetSize())
        self.SetMinSize(self.GetSize())

        self.Bind(wx.EVT_CLOSE, self.OnExit)

        self.thread = Thread(target=self.run_tasks)
        self.thread.setDaemon(True)
        self.thread.start()

    def OnExit(self, evt):
        self.STOP_TASK = True
        self.DestroyLater()

    def update_process(self, msg, tip):
        if self.file_lab is not None:
            self.file_lab.SetLabel("")
        if self.file_msg is not None:
            self.file_msg.SetLabel("")

        self.file_lab = TransparentStaticText(self.back, label=msg)
        self.file_lab.SetFont(
            wx.Font(
                12,
                wx.FONTFAMILY_MODERN,
                wx.FONTSTYLE_NORMAL,
                wx.FONTWEIGHT_LIGHT,
                False,
                "Microsoft YaHei",
            )
        )
        self.file_lab.SetForegroundColour((0, 0, 255))
        self.file_lab.SetPosition((350, 50))

        self.file_msg = TransparentStaticText(self.back, label=tip)
        self.file_msg.SetFont(
            wx.Font(
                12,
                wx.FONTFAMILY_MODERN,
                wx.FONTSTYLE_NORMAL,
                wx.FONTWEIGHT_LIGHT,
                False,
                "Microsoft YaHei",
            )
        )
        self.file_msg.SetForegroundColour((255, 255, 255))
        self.file_msg.SetPosition((100, 100))

    def run_tasks(self):
        with ThreadPoolExecutor(max_workers=8) as pool:
            tasks = [
                pool.submit(WordFreqCount().fold_task, (file_, self.keywords))
                for file_ in self.files
            ]
            for i, t in enumerate(as_completed(tasks)):
                if self.STOP_TASK:
                    return
                status, msg, filename, data = t.result()
                self.task_queue.append([status, msg, filename, data])
                wx.CallAfter(self.gauge.SetValue, i + 1)
                msg = f"{round((i + 1) / len(self.files) * 100, 1)}%"
                wx.CallAfter(self.gauge.SetValue, i + 1)
                wx.CallAfter(self.file_lab.SetLabel, msg)
                wx.CallAfter(
                    self.file_msg.SetLabel, f"解析文件：{os.path.basename(filename)}"
                )
        wx.CallAfter(self.Destroy)


if __name__ == "__main__":
    app = wx.App()
    with WordsFreqCountUI(None) as dlg:
        dlg.ShowModal()
    app.MainLoop()
